#任务计划者
'''现在长进了吗，我不觉得的，能困住我的事情，还有好多，变得更弱了吗，我也不觉得，毕竟那些经过，得算些什么'''
from pydantic import BaseModel,Field
from typing import List
from model_utils import getLLM

# {'task_list': ['查找2024年法国奥运会10米跳水冠军姓名', '搜索该冠军父亲的信息']}
# {'task_list': ['查找2024法国奥运会10米跳水项目的冠军', '获取该冠军运动员的个人信息', '确定该运动员的父亲是誰']}
# {'task_list': ['查找2024法国奥运会10米跳水冠军的信息', '获取该冠军运动员父亲的名字']}
class Plan(BaseModel):
    task_list:List[str] = Field(description="子任务列表")

_system_prompt_template = '''
你是一个优秀的规划师，根据查询目标，规划子目标
输出
{output_format}
'''

_human_prompt_template = '''
查询目标:
{query}
'''

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
class Planner:
    def __init__(self,llm):
        _prompt = ChatPromptTemplate.from_messages([
            ("system",_system_prompt_template),
            ("human",_human_prompt_template)
        ])
        _parser = JsonOutputParser(pydantic_object=Plan)
        # KeyError: "Input to ChatPromptTemplate is missing variables {'output_format'}.  Expected: ['output_format',
        # _prompt.partial(output_format=_parser.get_format_instructions())
        print(_parser.get_format_instructions())
        _prompt = _prompt.partial(output_format=_parser.get_format_instructions())
        print(_prompt)
        self.chain = _prompt | llm | _parser
    def __call__(self,state):
        return self.chain.invoke(state)

if __name__ == '__main__':
    llm = getLLM()
    _planner = Planner(llm)
    _rt = _planner({"query":"2024年法国巴黎奥运会女子10米跳水冠军的家乡在哪里"})
    print(_rt)